# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/10 9:58 
# @Author  : shutao
# @FileName: script_colloction.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

from loguru import logger
import os.path as osp
import numpy as np
import random
import cv2 as cv
import os


def create_list_of_dataset(root_path="E:\\work\\Dataset\\house_dataset",
                           image_folder_name="images",
                           label_forder_name="seg_images",
                           suffix=(".png", ".json"),
                           train_val_test=None):
    """
    Create and split the data list

    How to use this function:
        Pass in the data root directory, as well as the image folder and label folder under the data set root
        directory. The images in the image folder need to correspond to the label file name under the label folder
        one-to-one, and the corresponding suffix can be in `suffix` parameter settings. This method will generate
        three files named `train.list`、 `val.list`、 and `test.list`. The last, it will be save them to the root
        directory of the dataset.(note: the default split ratio is 6:1:3)
    Args:
        root_path: Data sets root path.
        image_folder_name: Image folder under the root directory of the dataset.
        label_forder_name: Label folder under the root directory of the data set.
        suffix: Suffixes corresponding to images and labels.
        train_val_test: Train val test split ratio.

    Returns:
        None
    """
    if train_val_test is None:
        train_val_test = [6, 1, 3]
    assert sum(train_val_test) == 10, "The sum of the division ratios must be 10!!!"

    # Create data list file
    train_list = open(osp.join(root_path, "train.list"), "w")
    val_list = open(osp.join(root_path, "val.list"), "w")
    test_list = open(osp.join(root_path, "test.list"), "w")

    number_of_copies = 10
    image_list = os.listdir(osp.join(root_path, image_folder_name))
    # Shuffle the data list
    random.shuffle(image_list)
    for idx, _file in enumerate(image_list):
        img_path = osp.join(image_folder_name, _file)
        lab_path = osp.join(label_forder_name, _file.replace(suffix[0], suffix[1]))

        if not osp.exists(osp.join(root_path, lab_path)):
            "The target file does not exist, pass it：{}".format(lab_path)
            continue
        # Create a row corresponding to the data list
        write_line = "{}\t{}\n".format(img_path, lab_path)
        logger.info(write_line, end="")

        # Write the data list file according to the split ratio
        if number_of_copies > 10 - train_val_test[0]:
            train_list.writelines(write_line)
        elif number_of_copies > 10 - sum(train_val_test[:2]):
            val_list.writelines(write_line)
        elif number_of_copies > 10 - sum(train_val_test):
            test_list.writelines(write_line)
        else:
            number_of_copies = 11
        number_of_copies -= 1

    # End of operation, close the data file
    train_list.close()
    val_list.close()
    test_list.close()
    logger.info("Create and split the data list complete!")


def parse_img():
    img_dir = "DSC_0017.jpg"
    img = cv.imread(img_dir)
    for c in range(img.shape[0]):
        for b in range(img.shape[1]):
            tmp = img[c, b]
            if tmp[0] <= 100 and tmp[1] <= 100 and tmp[2] >= 180:
                img[c, b] = [255, 255, 255]
    cv.imwrite("res.jpg", img)
    pass


if __name__ == '__main__':
    # create_list_of_dataset(root_path="E:\\work\\Dataset\\paper_document_dataset",
    #                        image_folder_name="Image",
    #                        label_forder_name="Annotations",
    #                        suffix=(".jpg", ".txt"))
    parse_img()
    pass
